const tf = require("@tensorflow/tfjs-node"); // TensorflowJs-Node
const getTrainData = require("./train"); // 方法：获取模型的训练数据

const main = async () => {
    /* ******************** 1.获取模型的训练数据 ******************** */

    /* DATA_DIR 训练数据图片目录(相对路径)，后期根据需求变化，可快速更改目录下的图片数据
     * SAVE_DIR 程序输出文件目录(相对路径)，自定义目录，程序中需要输出文件以记录某些数据
     * getTrainData 自定义方法，获取模型的训练数据，并处理成能投喂给模型的格式和类型 */
    const DATA_DIR = "images/train";
    const SAVE_DIR = "model-save";
    const {
        inputTs,  // 图片资源的图像张量Tensor(模型输入)
        outputTs, // 图片资源的类别标记Tensor(模型输出)
        category  // 图片资源的类别名称数组
    } = await getTrainData(DATA_DIR, SAVE_DIR);
    console.log(inputTs, outputTs, category);

    /* ******************** 2.定义垃圾的分类模型 ******************** */

    /* MobileNet是谷歌推出的轻量级神经网络，能够在移动终端实现众多应用，如目标检测及分类和人脸识别等
     * 相比大型神经网络，MobileNet拥有更小体积、更少计算量和更高精度，在轻量级神经网络中占据极大优势
     * MobileNet凭借网络结构优化，能够在更少的参数和计算量情况下，其网络精度反超了部分大型神经网络 */
    const MOBILE_NET = "http://ai-sample.oss-cn-hangzhou.aliyuncs.com/pipcook/models/mobilenet/web_model/model.json";

    /* tf.loadLayersModel() 从服务器加载由层对象组成的模型(即配置文件)，从而获取模型 */
    const MobileNet = await tf.loadLayersModel(MOBILE_NET);
    MobileNet.summary(); // 以文本摘要的形式，打印出模型的图层结构

    /* MobileNet.layers 模型图层，使用此方式查看模型的图层结构，能够比较直观地呈现出数据的格式类型
     * 终端打印数据显示，整体数据类型是数组(Array)，数组内的元素是具有多个属性的复杂对象(Object) */
    console.log(MobileNet.layers); // 模型的图层结构，是由对象为元素组成的数组

    /* 迁移训练，主要是指在复杂的深度学习模型中，通常有数百万的参数，若是从头完整地训练模型，则需要大量的计算资源
     * 现有较成熟完善的模型，或是已在相关任务中训练过的模型，复用其中可利用的部分，便可降低对大量计算资源的需要 */
    /* MobileNet模型，能够对非常多(数千种)的物体进行训练和分类，不仅是垃圾，还有动植物等其他各种类别，可以拿来复用到垃圾分类模型中
     * 复用MobileNet模型的前面部分，省去繁琐复杂的训练过程，在最后分类时，再自行定义成需要输出的垃圾类别(截断模型+双层神经网络) */

    /* 查看模型的图层结构数据，获取需要截断的层数索引值，用以截断模型(前面部分层直接复用，后面按需自行定义)
     * 据终端数据显示，需要截断的模型层信息：layer.name="conv_pw_13_relu"，index=86(总体第87层) */
    MobileNet.layers.map((layer, index) => { console.log(layer.name, index); });

    /* tf.sequential() 创建空的顺序/序列模型(即上层的输出是下层的输入) */
    /* Model.add() 在模型层堆栈顶部添加图层实例，此处是部分复用已现成的模型，追加到新创建的模型中
     * 根据上述打印代码查看得知，复用MobileNet模型的前87层(层数索引值0~86)，图片可无需执行训练 */
    const RubbishModel = tf.sequential();
    for (let i = 0; i <= 86; i++) {
        const layer = MobileNet.layers[i];
        layer.trainable = false; // 跳过训练，直接复用
        RubbishModel.add(layer); // 将图层结构追加到模型
    }

    /* tf.layers.flatten() 展平输入，将多维输入一维化，将高维度数据展平，常用在从卷积层到全连接层的过渡，方便后续模型做分类
     * tf.layers.dense() 定义全连接层，此处是在模型中添加全连接层，且添加两次(units:神经元个数，activation:激活函数) */
    /* 为提升CNN(卷积神经网络)性能，激活函数通常使用ReLU(线性整流函数/修正线性单元)，此处是指非线性函数，用来解决非线性问题
     * 最后的全连接层的输出值，被传递给一个输出，可以使用Softmax(归一化指数函数)逻辑回归进行分类，此处用来解决多分类问题 */
    /* 全连接层中的每个神经元，都与上层的所有神经元相连，将前面层提取到的特征综合起来，由于全相连特性，通常全连接层的参数也最多 */
    RubbishModel.add(tf.layers.flatten());
    RubbishModel.add(tf.layers.dense({ units: 10, activation: "relu" }));                 // 双层神经网络：隐藏层
    RubbishModel.add(tf.layers.dense({ units: category.length, activation: "softmax" })); // 双层神经网络：输出层

    /* ******************** 3.训练模型并将训练后的参数保存成文件 ******************** */

    /* Model.compile() 模型编译(配置训练方法，设置训练时的评测标准)，编译为模型配备损失函数(loss)、优化器(optimizer)及度量单位(metrics)等 */
    /* sparseCategoricalCrossentropy 某种分类相关的损失函数，图片分类是有固定损失函数的，因模型输出是类别目录下标(即数组下标0123)而选择该函数
     * 优化器，在深度学习的反向传播过程中，指引损失函数的各参数，往正确的方向上更新合适的大小，更新后的各参数，能够让损失函数值不断地逼近全局最小值 */
    RubbishModel.compile({
        loss: "sparseCategoricalCrossentropy", // 损失函数
        optimizer: tf.train.adam(), // 某种优化器，具有调配学习率和动画演示的功能
        metrics: ["acc"] // 准确度(Accuracy)的度量
    });

    /* Model.fit() 执行模型训练，拟合(使模型参数尽可能地拟合训练数据)，未编译(未调用过compile()方法)的模型不能执行fit()函数
     * epochs 迭代次数(即训练轮数)，模型将所有图片完整地学习训练一遍算作一轮，并非轮数越多越好，过多可能会产生过度拟合等问题 */
    await RubbishModel.fit(inputTs, outputTs, { epochs: 20 });

    /* 模型训练完成后，内部参数都会调整好，若不将这些配置参数保存下来，下次使用模型时还需重新训练，因此将训练后的参数保存成文件并导出
     * 参见目录model-save，生成文件model.json(模型原信息，如层结构配置等)和weights.bin(权重和参数等，机器自行调整的参数等) */
    /* process.cwd() 返回NodeJs进程的当前工作目录(绝对路径) /Users/vyingg/Documents/WorkSpace/Rubbish-Classify */
    await RubbishModel.save(`file://${process.cwd()}/${SAVE_DIR}`); // Rubbish-Classify/model-save
};
main(); // 执行主函数
